from concurrent.futures import (
    ProcessPoolExecutor,
    ThreadPoolExecutor
)
import aiomultiprocess
import functools
import asyncio
from typing import Callable
from enum import Enum, IntEnum
from typing import Optional
import os

from common.log import log
from config.base import get_settings

settings = get_settings()

class AsyncThreadPool:
    """
    异步线程池工具类，用于执行同步阻塞任务。
    """
    
    def __init__(self, max_workers: Optional[int] = None):
        """
        初始化线程池。

        :param max_workers: 线程池的最大工作线程数，默认为 `os.cpu_count()` 的值。
        """
        self.max_workers = max_workers or settings.thread.WORKERS
        self.executor = ThreadPoolExecutor(max_workers=self.max_workers)

    async def run_in_thread(self, func: Callable, timeout: Optional[float] = None, *args, **kwargs):
        """
        在线程池中异步执行一个函数。

        :param func: 要执行的函数。
        :param args: 函数的位置参数。
        :param kwargs: 函数的关键字参数。
        :return: 函数的返回值。
        """
        loop = asyncio.get_running_loop()
        wrapped_func = functools.partial(func, *args, **kwargs)
        try:
            # 添加超时控制
            if timeout is not None:
                return await asyncio.wait_for(loop.run_in_executor(self.executor, wrapped_func), timeout=timeout)
            else:
                return await loop.run_in_executor(self.executor, wrapped_func)
        except asyncio.TimeoutError:
            log.warning(f"任务超时: {func.__name__}")
            raise
        except Exception as e:
            log.error(f"任务执行失败: {e}")
            raise

    async def shutdown(self, wait: bool=True):
        """
        关闭线程池。

        :param wait: 是否等待所有任务完成后再关闭线程池，默认为 `True`。
        """
        await asyncio.to_thread(self.executor.shutdown, wait)